#!/usr/bin/python3
# -*- coding:utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================

import os
import torch
import numpy as np

def channel_mixing(x: torch.Tensor, h0: torch.Tensor, x_k: torch.Tensor, ffn_key: torch.Tensor, ffn_value: torch.Tensor) -> torch.Tensor:
    """
    通道混合函数。

    Args:
        x (torch.Tensor): 输入张量，形状为[B, T, C]。
        h0 (torch.Tensor): 输入张量(state状态)，形状为[B, 1, C]。
        x_k (torch.Tensor): 输入张量(模型权重），形状为[1, 1, C]。
        ffn_key (torch.Tensor): 输入张量(模型权重），形状为[4*C, C]。
        ffn_value (torch.Tensor): 输入张量(模型权重），形状为[C, 4*C]。
    Returns:
        out (torch.Tensor): 输出张量，形状为[B, T, C]。
        ht (torch.Tensor): 输出张量(state状态),形状为[B, 1, C]。
    """
    batch_size, seq_length = x.shape[0], x.shape[1]
    if seq_length == 1:
        sx = h0 - x
        ht = x
    else:
        h0 = h0.view(batch_size,1,-1)
        h0 = torch.cat([h0, x[:, :-1, :]], dim=1)
        sx = (h0 - x)
        ht = x[:, -1, :]
    xk = x + sx * x_k
    k = torch.relu(xk @ ffn_key.T).pow(2)
    return k @ ffn_value.T, ht

def gen_golden_data_simple():
    dtype = np.float16

    my_case = os.environ["CHANNEL_MIXING_CASE"]

    if my_case == "1": 
        B = 1
        T = 1
    elif my_case == "2": 
        B = 1
        T = 8
    elif my_case == "3": 
        B = 1
        T = 64
    elif my_case == "4": 
        B = 1
        T = 1024
    elif my_case == "5": 
        B = 1
        T = 8192
    elif my_case == "6": 
        B = 7
        T = 77
    elif my_case == "7": 
        B = 8
        T = 32
    elif my_case == "8": 
        B = 9
        T = 256
    elif my_case == "9": 
        B = 5
        T = 80
    elif my_case == "10": 
        B = 6
        T = 128
    elif my_case == "11": 
        B = 1
        T = 123
    elif my_case == "12": 
        B = 100
        T = 16
    elif my_case == "13": 
        B = 200
        T = 64
    elif my_case == "14": 
        B = 11
        T = 127
    elif my_case == "15": 
        B = 3
        T = 2048
    elif my_case == "0": 
        B = 1
        T = 80
    
    C = 2560

    x = np.random.uniform(-0.1, 0.1, [B, T, C]).astype(dtype)
    h0 = np.random.uniform(-0.1, 0.1, [B, 1, C]).astype(dtype)
    xk = np.random.uniform(-0.1, 0.1, [1, 1, C]).astype(dtype)
    kw = np.random.uniform(-0.1, 0.1, [4 * C, C]).astype(dtype)
    vw = np.random.uniform(-0.1, 0.1, [C, 4 * C]).astype(dtype)

    out, ht = channel_mixing(torch.from_numpy(x), torch.from_numpy(h0), torch.from_numpy(xk), torch.from_numpy(kw), torch.from_numpy(vw))

    os.system("mkdir -p input")
    x.tofile("./input/input_x.bin")
    h0.tofile("./input/input_h0.bin")
    xk.tofile("./input/input_xk.bin")
    kw.tofile("./input/input_kw.bin")
    vw.tofile("./input/input_vw.bin")

    os.system("mkdir -p output")
    out.numpy().tofile("./output/out.bin")
    ht.numpy().tofile("./output/ht.bin")

if __name__ == "__main__":
    gen_golden_data_simple()

